#!/usr/bin/python3
# _*_ coding: utf-8 _*_
#
# Copyright (C) 2024 - 2024 heihieyouheihei, Inc. All Rights Reserved 
#
# @Time    : 2024/4/14 10:04
# @Author  : Alchemist_Niu
# @File    : main_ui.py
# @IDE     : PyCharm
from PyQt5.QtWidgets import *
from PyQt5 import QtGui
from PyQt5.QtCore import *
from win.yolov5 import Ui_MainWindow
import sys,os,time, datetime
from detect_plate import *
from PyQt5.QtGui import QImage, QPixmap

class myMainWindow(Ui_MainWindow, QMainWindow):
    def __init__(self):
        super(myMainWindow, self).__init__()
        self.setupUi(self)
        self.timer_camera = QTimer()
        self.flag = True
        self.flag_1 = False
        self.file_type = None
        self.camera_opened = False

        self.PB1.clicked.connect(self.OpenFile)         #绑定按键1：打开文件
        self.PB2.clicked.connect(self.OpenCamera)        #打开/关闭摄像头
        self.PB3.clicked.connect(self.RunOrContinue)     #开始/暂停检测
        self.PB4.clicked.connect(self.Stop)               #结束检测

        self.det_thread = DetThread()
        self.det_thread.send_img.connect(lambda x: self.show_image(x, self.label_1))
        self.det_thread.send_result.connect(lambda x: self.statistic_msg(x))

    def OpenFile(self):
        self.name, _ = QFileDialog.getOpenFileName(self, "打开文件", "./imgs/",
                                            "All Files(*);;*.jpg;;*.png;;*.mp4;;*.avi")
        _, self.file_type = os.path.splitext(self.name)  # 将文件路径分离文件名与扩展名
        if self.name:  # 对读取到的文件类型进行判断
            if self.name.endswith(".jpg") or self.name.endswith(".png"):  # 如果文件为.jpg格式或.png格式
                ori_img = cv2.imread(self.name)  # 使用cv2函数读取该图片
                self.show_image(ori_img, self.label_1)  # 将该图片显示到GUI界面上

            elif self.name.endswith(".mp4") or self.name.endswith(".avi"):  # 如果文件为.mp4格式或.avi格式
                cap = cv2.VideoCapture(self.name)  # 使用cv2.VideoCapture函数读取视频
                ret, frame = cap.read()  # 读取该视频的第一帧
                self.show_image(frame, self.label_1)  # 读取到的第一帧图片显示到GUI界面上

        self.det_thread.path = self.name

    def OpenCamera(self):
        if not self.camera_opened:
            self.file_type = "camera"
            self.camera_opened = True
            self.video = cv2.VideoCapture(0)  # 读取视频文件
            fps = int(self.video.get(cv2.CAP_PROP_FPS))  # 获取视频的帧率
            while self.video.isOpened():
                ret, frame = self.video.read()
                if frame is None:
                    break
                if ret:
                    self.show_image(frame, label=self.label_1)
                    cv2.waitKey(1000 // fps)

            self.video.release()
            cv2.destroyAllWindows()

        else:  # 如果摄像头已打开，则关闭摄像头
            self.video.release()
            if self.timer_camera.isActive():
                self.timer_camera.stop()
            cv2.destroyAllWindows()
            self.camera_opened = False  # 更新摄像头状态

    def RunOrContinue(self):
        self.det_thread.jump_out = False
        if not self.det_thread.isRunning():  # 如果检测线程没有在运行
            if self.file_type is not None and self.file_type in ".jpg .png":
                self.det_thread.start() #开启检测线程
                self.statistic_msg('开始检测>>>')  # 发送结束信息

            elif self.file_type is not None and self.file_type in ".mp4 .avi":
                self.det_thread.start()  # 开启检测线程
                self.statistic_msg('开始检测>>>')  # 发送结束信息
            elif self.file_type is not None and self.file_type == "camera":
                self.flag_1 = True
                self.det_thread.start()  # 开启检测线程
                self.statistic_msg('开始检测>>>')  # 发送结束信息

            else:
                QMessageBox.warning(self, "提示", "未选择文件！！")

        else:
            QMessageBox.warning(self,"提示","检测还未结束！")

    def Stop(self):
        self.det_thread.jump_out = True
        self.flag_1 = False
        self.file_type = None
        self.statistic_msg('结束检测>>>')  # 发送结束信息

    '''
    发送消息
    '''
    def statistic_msg(self, msg):
        self.textBrowser.append(msg) #将接收到的信息内容显示到statistic_label标签上

    '''
    函数功能：将图片信息流显示到GUI界面的指定标签上
    '''
    @staticmethod
    def show_image(img_src, label):
        try:
            ih, iw, _ = img_src.shape
            w = label.geometry().width()
            h = label.geometry().height()
            # keep original aspect ratio
            if iw / w > ih / h:
                scal = w / iw
                nw = w
                nh = int(scal * ih)
                img_src_ = cv2.resize(img_src, (nw, nh))

            else:
                scal = h / ih
                nw = int(scal * iw)
                nh = h
                img_src_ = cv2.resize(img_src, (nw, nh))
            frame = cv2.cvtColor(img_src_, cv2.COLOR_BGR2RGB)
            img = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[2] * frame.shape[1],
                         QImage.Format_RGB888)
            label.setPixmap(QPixmap.fromImage(img))

        except Exception as e:
            print(repr(e))

class DetThread(QThread):
    send_img = pyqtSignal(np.ndarray)  # 自定义信号，用于传输原始图像信号
    send_result = pyqtSignal(str) #发送结果信号

    def __init__(self):
        super().__init__()
        self.path = None
        self.timer_video = QTimer()
        self.jump_out = False  # jump out of the loop，跳出循环
        self.plate_result = None

        self.detect_model = load_model(opt.detect_model, device)  # 初始化yolov8识别模型
        self.plate_rec_model = init_model(device, opt.rec_model, is_color=True)  # 初始化识别模型


    def run(self):
        # 算参数量
        total = sum(p.numel() for p in self.detect_model.parameters())
        total_1 = sum(p.numel() for p in self.plate_rec_model.parameters())
        print("yolov8 detect params: %.2fM,rec params: %.2fM" % (total / 1e6, total_1 / 1e6))
        self.send_result.emit("yolov8 detect params: %.2fM,rec params: %.2fM" % (total / 1e6, total_1 / 1e6))
        self.detect_model.eval()
        if self.path is not None:
            self.send_result.emit("检测的文件为："+ self.path)
            _, self.file_type = os.path.splitext(self.path)  # 将文件路径分离文件名与扩展名
            if self.file_type in ".jpg .png":#如果文件为.jpg格式或.png格式
                try:
                    img = cv_imread(self.path)  # cv2读取图片
                    if img.shape[-1] == 4:
                        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)  # 图像格式转换，RGB转BGR
                    # img_ori = copy.deepcopy(img)
                    result_list = detect_Recognition_plate(self.detect_model, img, device, self.plate_rec_model, opt.img_size, is_color=opt.is_color)
                    ori_img, result = draw_result(img, result_list)  # 原图上绘制检测
                    self.send_img.emit(ori_img) #将检测结果发送出去
                    self.send_result.emit("检测结果："+ result)
                    img_name = os.path.basename(self.path)  # 获取图片名称
                    save_img_path = os.path.join(opt.output, img_name)  # 路径拼接，将检测结果保存至新的文件夹下
                    cv2.imwrite(save_img_path, ori_img)  # 保存检测结果

                    self.path = None

                except Exception as e:
                    print(e)

            elif self.file_type in ".mp4 .avi":#如果文件为.jpg格式或.png格式
                try:
                    video = cv2.VideoCapture(self.path)  # 读取视频文件
                    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取视频的宽度
                    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取视频的高度
                    fps = int(video.get(cv2.CAP_PROP_FPS))  # 获取视频的帧率

                    img_name = os.path.basename(self.path)
                    save_img_path = os.path.join(opt.output, img_name)
                    videoWriter = cv2.VideoWriter(save_img_path, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps,
                                                  (width, height))
                    while video.isOpened():
                        if self.jump_out:
                            video.release()
                            if self.timer_video.isActive():
                                self.timer_video.stop()
                            cv2.destroyAllWindows()
                            break
                        ret, frame = video.read()
                        # frame_ori = copy.deepcopy(frame)
                        if frame is None:
                            break
                        if ret:
                            dict_list = detect_Recognition_plate(self.detect_model, frame, device, self.plate_rec_model, opt.img_size, is_color=opt.is_color)
                            video_result, plate_result_1 = draw_result(frame, dict_list)
                            self.send_img.emit(video_result)
                            if plate_result_1 != self.plate_result: #对视频中的同一个车牌限制输出
                                self.send_result.emit("检测结果："+ plate_result_1)
                                self.plate_result = plate_result_1
                            videoWriter.write(video_result)

                    video.release()
                    cv2.destroyAllWindows()
                    self.path = None

                except Exception as e:
                    print(e)

        else:
            self.send_result.emit("检测的文件为：摄像头")
            video = cv2.VideoCapture(0)  # 读取视频文件
            width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取视频的宽度
            height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取视频的高度
            fps = int(video.get(cv2.CAP_PROP_FPS))  # 获取视频的帧率

            current_time = datetime.datetime.now()
            formatted_time = current_time.strftime("%Y%m%d_%H%M%S")
            file_name = f"data_{formatted_time}.mp4"
            save_img_path = os.path.join(opt.output, file_name)
            videoWriter = cv2.VideoWriter(save_img_path, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps,
                                          (width, height))
            while video.isOpened():
                if self.jump_out:
                    video.release()
                    if self.timer_video.isActive():
                        self.timer_video.stop()
                    cv2.destroyAllWindows()
                    break
                ret, frame = video.read()
                # frame_ori = copy.deepcopy(frame)
                if frame is None:
                    break
                if ret:
                    dict_list = detect_Recognition_plate(self.detect_model, frame, device, self.plate_rec_model, opt.img_size, is_color=opt.is_color)
                    video_result, plate_result_1 = draw_result(frame, dict_list)
                    self.send_img.emit(video_result)
                    if plate_result_1 != self.plate_result:  # 对视频中的同一个车牌限制输出
                        self.send_result.emit("检测结果：" + plate_result_1)
                        self.plate_result = plate_result_1
                    videoWriter.write(video_result)

            video.release()
            cv2.destroyAllWindows()
            self.path = None


# QCoreApplication.setAttribute(Qt.AA_EnableHighDpiScaling)
if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--detect_model', nargs='+', type=str, default='weights/plate_detect.pt', help='model.pt path(s)')  #检测模型
    parser.add_argument('--rec_model', type=str, default='weights/plate_rec_color.pth', help='model.pt path(s)')#车牌识别+颜色识别模型
    parser.add_argument('--is_color',type=bool,default=True,help='plate color')      #是否识别颜色
    parser.add_argument('--image_path', type=str, default='imgs', help='source')     #图片路径
    parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')  #网络输入图片大小
    parser.add_argument('--output', type=str, default='result', help='source')               #图片结果保存的位置
    parser.add_argument('--video', type=str, default='imgs/plate_test.mp4', help='source')                       #视频的路径
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")                     #使用gpu还是cpu进行识别
    clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
    opt = parser.parse_args()
    save_path = opt.output
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    app = QApplication(sys.argv)
    vieo_gui = myMainWindow()
    vieo_gui.show()
    sys.exit(app.exec_())